package com.tang.state;

import com.tang.bean.WaterSensor;
import com.tang.functions.WaterSensorMapFunction;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.scala.typeutils.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.BroadcastConnectedStream;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;

import java.util.Optional;

/**
 * 水位超过指定的阈值发送告警，阈值可以动态修改。
 *
 * @author tang
 * @since 2023/7/11 21:54
 */
public class OperatorBroadcastStateDemo {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(new Configuration());
        env.setParallelism(2);

        // 数据流
        SingleOutputStreamOperator<WaterSensor> sensorDataSource = env
                .socketTextStream("192.168.70.141", 7777)
                .map(new WaterSensorMapFunction());

        // 配置流（用来广播配置）
        DataStreamSource<String> configDataSource = env
                .socketTextStream("192.168.70.141", 8888);

        // 1.将配置流广播
        MapStateDescriptor<String, Integer> broadcastMapState =
                new MapStateDescriptor<>("broadcast-state", Types.STRING(),
                        Types.INT());
        BroadcastStream<String> configBroadcastStream = configDataSource.broadcast(broadcastMapState);

        // 2.把数据流和广播后的配置类connect
        BroadcastConnectedStream<WaterSensor, String> sensorBcs = sensorDataSource
                .connect(configBroadcastStream);

        // 3.调用process
        sensorBcs.process(new BroadcastProcessFunction<WaterSensor, String, String>() {

                    /**
                     * 对于(非广播)数据流中的每个元素调用此方法。
                     * 该函数可以使用Collector参数输出零个或多个元素，查询当前处理/事件时间，
                     * 还可以查询和更新本地键控状态。最后，它对广播状态具有只读访问权限。
                     * 上下文仅在调用此方法期间有效，不要存储它。
                     * 
                     * @param value 流数据
                     * @param ctx 上下文
                     * @param out 输出
                     * @throws Exception -
                     */
                    @Override
                    public void processElement(WaterSensor value, BroadcastProcessFunction<WaterSensor, String, String>.ReadOnlyContext ctx, Collector<String> out) throws Exception {
                        System.out.println("这里是非广播流数据哦哦哦哦哦");
                        // 你看，顾名思义，这是只读的
                        ReadOnlyBroadcastState<String, Integer> broadcastState = ctx.getBroadcastState(broadcastMapState);
                        // 判断广播里是否有数据，因为刚刚启动时，可能数据流比广播流的数据先来
                        Integer threshold = Optional.ofNullable(broadcastState.get("threshold"))
                                .orElse(0);
                        if (value.getVc() > threshold) {
                            out.collect(value + ", 水位超过指定阈值：" + threshold + "!!!!");
                        }
                    }

                    /**
                     * 对于广播流中的每个元素调用此方法。
                     * 该函数可以使用Collector参数输出零个或多个元素，
                     * 查询当前处理/事件时间，还可以查询和更新内部广播状态。
                     * 这些可以通过提供的BroadcastProcessFunction.Context来完成。
                     * 上下文仅在调用此方法期间有效，不要存储它。
                     * 
                     * @param value element
                     * @param ctx 上下文
                     * @param out 输出
                     * @throws Exception -
                     */
                    @Override
                    public void processBroadcastElement(String value, BroadcastProcessFunction<WaterSensor, String, String>.Context ctx, Collector<String> out) throws Exception {
                        // 4. 通过上下文获取广播状态，并且往里写数据
                        System.out.println("这是一条广播流的数据噢噢噢噢");
                        BroadcastState<String, Integer> broadcastState = ctx.getBroadcastState(broadcastMapState);
                        broadcastState.put("threshold", Integer.valueOf(value));
                        System.out.println("重新指定个阈值为：" + value);
                    }

                })
                .print();

        env.execute();
    }

}
